/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#include "mstep.h"
#include "qfunc.h"

#include <assert.h>
#include <gsl/gsl_multimin.h>
#include <stdlib.h>
#include <string.h>

static void
to_x  (const double*    alpha,
       unsigned int     n_workers,
       unsigned int     n_labels,
       const double*    logbeta,
       unsigned int     n_images,
       gsl_vector*      x)
{
  memcpy (x->data, alpha, n_workers * n_labels * n_labels * sizeof (alpha[0]));
  memcpy (x->data + n_workers * n_labels * n_labels,
          logbeta,
          n_images * sizeof (logbeta[0]));
}

static void
from_x (double*                 alpha,
        unsigned int            n_workers,
        unsigned int            n_labels,
        double*                 logbeta,
        unsigned int            n_images,
        const gsl_vector*       x)
{
  memcpy (alpha, x->data, n_workers * n_labels * n_labels * sizeof (alpha[0]));
  memcpy (logbeta, 
          x->data + n_workers * n_labels * n_labels,
          n_images * sizeof (logbeta[0]));
}

typedef struct _QfuncParams QfuncParams;
struct _QfuncParams 
{
  unsigned int             n_workers;
  unsigned int             n_labels;
  unsigned int             n_images;
  const double*            pz;
  const double*            priorz;
  const Rating*            ratings;
  unsigned int             n_ratings;
  const Distribution*      prioralpha;
  const Distribution*      priorlogbeta;
  int                      symmetric;
};

static double
qfunc_f (const gsl_vector* x,
         void*             params)
{
  QfuncParams *qfunc_params = (QfuncParams *) params;
  unsigned int n_workers = qfunc_params->n_workers;
  unsigned int n_labels = qfunc_params->n_labels;

  return -qfunc (x->data,
                 n_workers,
                 n_labels,
                 x->data + n_workers * n_labels * n_labels,
                 qfunc_params->n_images,
                 qfunc_params->pz,
                 qfunc_params->priorz,
                 qfunc_params->ratings,
                 qfunc_params->n_ratings,
                 qfunc_params->prioralpha,
                 qfunc_params->priorlogbeta,
                 qfunc_params->symmetric);
}

static void
qfunc_df (const gsl_vector* x,
          void*             params,
          gsl_vector*       g)
{
  QfuncParams *qfunc_params = (QfuncParams *) params;
  unsigned int n_workers = qfunc_params->n_workers;
  unsigned int n_labels = qfunc_params->n_labels;
  unsigned int n_images = qfunc_params->n_images;

  dqfunc (g->data,
          g->data + n_workers * n_labels * n_labels,
          x->data,
          n_workers,
          n_labels,
          x->data + n_workers * n_labels * n_labels,
          n_images,
          qfunc_params->pz,
          qfunc_params->ratings,
          qfunc_params->n_ratings,
          qfunc_params->prioralpha,
          qfunc_params->priorlogbeta,
          qfunc_params->symmetric);

  for (unsigned int i = 0; i < n_workers * n_labels * n_labels + n_images; ++i)
    {
      g->data[i] *= -1;
    }
}

static void
qfunc_fdf (const gsl_vector* x,
           void*             params,
           double*           f,
           gsl_vector*       g)
{
  *f = qfunc_f (x, params);
  qfunc_df (x, params, g);
}

double
mstep (double*                  alpha,
       unsigned int             n_workers,
       unsigned int             n_labels,
       double*                  logbeta,
       unsigned int             n_images,
       const double*            pz,
       const double*            priorz,
       const Rating*            ratings,
       unsigned int             n_ratings,
       const Distribution*      prioralpha,
       const Distribution*      priorlogbeta,
       int                      symmetric,
       double                   tol)
{
  size_t n_dims = n_workers * n_labels * n_labels + n_images;
  gsl_multimin_fdfminimizer* s = 
    gsl_multimin_fdfminimizer_alloc (gsl_multimin_fdfminimizer_conjugate_pr,
                                     n_dims);
  assert (s != NULL);
  gsl_vector* x = gsl_vector_alloc (n_dims);
  assert (x != NULL);

  QfuncParams qfunc_params = { .n_workers = n_workers,
                               .n_labels = n_labels,
                               .n_images = n_images,
                               .pz = pz,
                               .priorz = priorz,
                               .ratings = ratings,
                               .n_ratings = n_ratings,
                               .prioralpha = prioralpha,
                               .priorlogbeta = priorlogbeta,
                               .symmetric = symmetric };

  gsl_multimin_function_fdf my_func = { .f = qfunc_f,
                                        .df = qfunc_df,
                                        .fdf = qfunc_fdf,
                                        .n = n_dims,
                                        .params = &qfunc_params };

  to_x (alpha, n_workers, n_labels, logbeta, n_images, x);

  gsl_multimin_fdfminimizer_set (s, &my_func, x, 0.01, 10 * tol);

  double lastf = 0;
  for (int iter = 0;
       ! gsl_multimin_fdfminimizer_iterate (s) &&
       gsl_multimin_test_gradient (s->gradient, tol) == GSL_CONTINUE &&
       (iter == 0 || lastf - s->f > 10 * tol) &&
       iter < 25;
       ++iter, 
       lastf = s->f)
    {
    }

  from_x (alpha, n_workers, n_labels, logbeta, n_images, s->x);

  double qmax = -s->f;

  gsl_vector_free (x);
  gsl_multimin_fdfminimizer_free (s);

  return qmax;
}
